package com.lin.linchatgpt.midjourney.service.impl;


import com.alibaba.fastjson.JSONObject;

import com.lin.linchatgpt.entity.dto.MidjourneyDTO;
import com.lin.linchatgpt.exception.LinException;
import com.lin.linchatgpt.midjourney.service.MidjourneyService;
import com.lin.linchatgpt.utils.ResultCode;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.concurrent.TimeUnit;

/**
 * Copyright(C),2022年-2022,霖霖
 *
 * @author 霖霖
 * @version 1.0
 * @date 2023/7/7 14:15
 * @Description
 */
@Service
@Slf4j
public class MidjourneyServiceImpl implements MidjourneyService {
    @Value("${midjourney.token}")
    private String token;
    @Value("${midjourney.host}")
    private String host;
    @Value("${midjourney.slow-host}")
    private String slowHost;
    @Value("${midjourney.slow-token}")
    private String slowToken;


    @Override
    public MidjourneyDTO getMidjourneyImage(String prompt, int modelType, int generateSpeed) {
        // TODO: 2023/7/7 后续改为从数据库中获取是否为慢速生成或者快速生成
        String url = getUrl(generateSpeed);

        MidjourneyDTO midjourneyDTO = null;
        // 异步回调
        switch (modelType) {
            case 1:
                midjourneyDTO = asyncCallback( prompt, url.toString());
                break;
            case 0:
//                streamGenerateImage(client, prompt, url.toString());
                break;
            default:
                throw new LinException(ResultCode.ERROR, "生成图片类型错误");
        }

        return midjourneyDTO;
    }

    private MidjourneyDTO asyncCallback( String prompt, String url) {
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("prompt", prompt);
        jsonObject.put("callback_url", "http://linchatgpt.linlin.online/lin_chat/generate_image/receive");
        return sendRequest(jsonObject, url);
    }

//    private MidjourneyDTO streamGenerateImage(OkHttpClient okHttpClient, String prompt, String url) {
//        MediaType mediaType = MediaType.Companion.parse("application/json; charset=utf-8");
//
//        JSONObject jsonObject = new JSONObject();
//        jsonObject.put("prompt", prompt);
//        RequestBody body = RequestBody.Companion.create(jsonObject.toJSONString(), mediaType);
//        Request request = new Request.Builder()
//                .url(url)
//                .post(body)
//                .addHeader("Content-Type", "application/json")
//                .addHeader("Accept", "application/x-ndjson")
//                .build();
//        MidjourneyDTO midjourneyDTO = null;
//        // 流式输出
//        okHttpClient.newCall(request).enqueue(new Callback() {
//            @Override
//            public void onFailure(Call call, java.io.IOException e) {
//                log.error("midjourney请求失败：{}", e.getMessage());
//            }
//
//            @Override
//            public void onResponse(Call call, Response response) throws java.io.IOException {
//                // 判断是否成功
//                if (!response.isSuccessful()) {
//                    log.error("midjourney请求失败：{}", response.message());
//                    throw new LinException(ResultCode.GENERATE_IMAGE_ERROR, response.toString());
//                }
//                if (response.body() != null) {
//                    String bodyJsonSrt = response.body().string();
//                    log.info("midjourney返回的数据为：{}", bodyJsonSrt);
//
//                    // 读取流式输入的数据
//                    try (BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(response.body().byteStream(), StandardCharsets.UTF_8))) {
//                        String line;
//                        while ((line = bufferedReader.readLine()) != null) {
//                            log.info("midjourney返回的数据流为：{}", line);
//                            midjourneyDTO = JSONObject.parseObject(line, MidjourneyDTO.class);
//                        }
//                    }
//                }
//            }
//        });
//    }


    @Override
    public MidjourneyDTO handlerImage(String action, String imageId) {
        String url = getUrl(1);

        JSONObject jsonObject = new JSONObject();
        jsonObject.put("action", action);
        jsonObject.put("image_id", imageId);
        jsonObject.put("callback_url", "http://linchat.linlin.online/lin_chat/generate_image/receive");
        return sendRequest(jsonObject, url);
    }

    private String getUrl(int generateSpeed){
        StringBuilder url = new StringBuilder();
        switch (generateSpeed) {
            case 0:
                url.append(host);
                url.append(token);
                break;
            case 1:
                url.append(slowHost);
                url.append(slowToken);
                break;
            default:
                throw new LinException(ResultCode.ERROR, "生成速度参数错误");
        }
        return url.toString();
    }

    private OkHttpClient getOkHttp(){
        return new OkHttpClient()
                .newBuilder()
                .connectTimeout(3 * 60, TimeUnit.SECONDS)
                .readTimeout(3 * 60, TimeUnit.SECONDS)
                .writeTimeout(3 * 60, TimeUnit.SECONDS)
                .build();
    }

    private MidjourneyDTO sendRequest(JSONObject jsonBody, String url){
        MediaType mediaType = MediaType.Companion.parse("application/json; charset=utf-8");
        RequestBody body = RequestBody.Companion.create(jsonBody.toJSONString(), mediaType);
        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .build();
        try (Response response = getOkHttp().newCall(request).execute()) {
            if (!response.isSuccessful()) {
                log.error("midjourney请求失败：{}", response.message());
                throw new LinException(ResultCode.GENERATE_IMAGE_ERROR, "midjourney请求失败");
            }
            if (response.body() != null) {
                String bodyJsonSrt = response.body().string();
                log.info("midjourney返回的数据为：{}", bodyJsonSrt);
                return JSONObject.parseObject(bodyJsonSrt, MidjourneyDTO.class);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
